import os
import json
import re
import logging
import platform
from collections import defaultdict

# 设置日志记录器，用于记录程序执行过程中的信息
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)  # 设置日志级别为DEBUG，记录所有级别的日志

# 文件日志处理器，将日志写入到../log.txt文件中
file_handler = logging.FileHandler('../log.txt', encoding='utf-8')
file_handler.setLevel(logging.DEBUG)  # 文件日志级别为DEBUG，记录所有级别的日志

# 控制台日志处理器，将日志输出到控制台
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)  # 控制台日志级别为INFO，只显示INFO及以上级别的日志

# 设置日志格式，包含时间、日志级别和消息内容
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 将处理器添加到日志记录器
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 检测操作系统类型
is_windows = platform.system() == "Windows"


# 路径处理函数，确保路径分隔符与操作系统兼容
def normalize_path(path):
    """
    根据当前操作系统规范化路径
    """
    if is_windows:
        return path.replace('/', '\\')
    else:
        return path.replace('\\', '/')


# 定义路径配置，根据操作系统类型使用不同的路径格式
if is_windows:
    # Windows路径
    json_root = r"\\172.28.10.19\车载全部素材\25.05.09目标识别测试素材汇总\晴天户外白天\test_frames"  # JSON文件所在的根目录
    txt_folder = r"\\172.28.10.19\ivqa\TestData\Zone_AI_Track\0522txt\inputtrack"  # TXT文件所在的文件夹
    output_file = r"D:\work\AI\20250523\track\test.txt"  # 结果输出文件路径
else:
    # Mac/Linux路径 - 使用SMB URL格式
    json_root = "/Users/only/车载全部素材/25.05.09目标识别测试素材汇总/晴天户外白天/test_frames"  # 假设已挂载到/Volumes
    txt_folder = "/Users/only/ivqa/TestData/Zone_AI_Track/0522txt/inputtrack"  # 假设已挂载到/Volumes
    output_file = "/Users/only/Desktop/工作文档/DH/path_code/trae_code/demo15/test.txt"  # 修复路径，移除重复的用户目录

    # 打印提示信息
    logger.info("在Mac上运行，请确保已挂载SMB共享")
    logger.info(f"如果共享未挂载，请使用mount_smbfs命令挂载")

# 规范化路径，确保路径分隔符与操作系统兼容
json_root = normalize_path(json_root)
txt_folder = normalize_path(txt_folder)
output_file = normalize_path(output_file)

# 创建输出文件夹，用于保存选择的结果和合并结果
txt_folder_name = os.path.basename(txt_folder.rstrip("/\\"))  # 获取txt文件夹的名称，兼容Windows和Mac的路径分隔符
output_dir = os.path.dirname(output_file)
if not os.path.exists(output_dir):
    os.makedirs(output_dir, exist_ok=True)  # 确保输出目录存在

# 创建合并结果的输出目录
merge_output_dir = os.path.join(output_dir, "merge_results")
os.makedirs(merge_output_dir, exist_ok=True)  # 确保合并结果目录存在
logger.info(f"合并结果将保存到: {merge_output_dir}")

# 创建选择结果的输出目录
choose_output_dir = os.path.join(os.path.dirname(output_file), txt_folder_name + "_choose")  # 创建选择结果的输出目录
os.makedirs(choose_output_dir, exist_ok=True)  # 确保目录存在，如果不存在则创建
choose_summary_path = os.path.join(choose_output_dir, txt_folder_name + "_choose.txt")  # 选择结果的汇总文件路径

# 类别映射字典，将数字ID映射到具体的目标类别名称
cls_map = {
    "0": "person",  # 人
    "1": "car",  # 车
    "2": "bike+person",  # 自行车+人
    "12": "tricycle",  # 三轮车
}

file_count = 0  # 初始化处理文件计数器
processed_results = defaultdict(list)  # 用于按目录分组存储处理结果

# 遍历JSON根目录下的所有文件
for root, _, files in os.walk(json_root):
    for file in files:
        # 只处理JSON文件
        if not file.endswith(".json"):
            continue

        file_count += 1  # 文件计数加1
        json_path = os.path.join(root, file)  # 构建完整的JSON文件路径
        json_path = normalize_path(json_path)  # 确保路径分隔符与操作系统兼容
        name_no_ext = os.path.splitext(file)[0]  # 获取不带扩展名的文件名
        logger.info(f"正在处理第 {file_count} 个文件: {file}")  # 记录当前处理的文件

        # 使用正则表达式从文件名中提取前缀和帧号
        # 文件名格式应为：前缀_六位数字帧号
        match = re.match(r"(.+)_([0-9]{6})$", name_no_ext)
        if not match:
            logger.warning(f"无法从文件名中提取帧号: {file}")  # 如果文件名格式不匹配，记录警告并跳过
            continue

        # 提取文件名中的前缀和帧号
        prefix, frame_str = match.groups()
        frame_num = int(frame_str.lstrip("0") or "0")  # 将帧号字符串转换为整数（去除前导0）
        txt_filename = prefix + ".txt"  # 构建对应的TXT文件名
        txt_path = os.path.join(txt_folder, txt_filename)  # 构建完整的TXT文件路径
        txt_path = normalize_path(txt_path)  # 确保路径分隔符与操作系统兼容
        logger.debug(f"匹配帧号: {frame_str}（转换为数字：{frame_num}）")
        logger.debug(f"尝试读取 TXT 文件: {txt_path}")

        extracted_objs = []  # 用于存储从TXT文件中提取的目标信息
        choose_lines = []  # 用于存储选择的行

        # 检查对应的TXT文件是否存在
        if os.path.exists(txt_path):
            try:
                # 读取TXT文件内容
                with open(txt_path, "r", encoding="utf-8") as txt_f:
                    found = False  # 初始化是否找到对应帧号的标志
                    for line in txt_f:
                        raw_line = line.strip()  # 去除行首尾的空白字符
                        if not raw_line:
                            continue  # 跳过空行

                        # 将行按空格分割，最多分成两部分（帧号和数据）
                        parts = raw_line.split(maxsplit=1)
                        try:
                            line_frame = int(parts[0])  # 尝试将第一部分转换为整数作为帧号
                            if line_frame == frame_num:  # 如果帧号匹配
                                found = True
                                choose_lines.append(raw_line)  # 将原始行添加到选择行列表

                                # 提取数据部分并确保格式正确（JSON数组格式）
                                data_str = parts[1] if len(parts) > 1 else "[]"
                                if not data_str.startswith("["):
                                    data_str = "[" + data_str
                                if not data_str.endswith("]"):
                                    data_str = data_str + "]"

                                # 解析JSON数据
                                objs = json.loads(data_str)

                                logger.debug(f"成功找到并解析帧号 {frame_num} 的 TXT 数据，共 {len(objs)} 项")

                                # 处理每个目标对象
                                for obj in objs:
                                    cls_id_str = str(obj.get("cls_id"))  # 获取类别ID
                                    if cls_id_str not in cls_map:
                                        logger.warning(f"未知的类别ID {cls_id_str}，跳过")
                                        continue
                                    # 如果类别ID在映射表中且包含矩形框信息
                                    if cls_id_str in cls_map and "rect" in obj:
                                        extracted_objs.append({
                                            "cls": cls_map[cls_id_str],  # 将类别ID映射为类别名称
                                            "bbox": obj["rect"]  # 保留边界框信息
                                        })
                                break  # 找到匹配帧号后退出循环
                        except ValueError:
                            logger.warning(f"TXT 行格式错误，跳过：{raw_line}")  # 记录格式错误的警告
                            continue
                    if not found:
                        logger.info(f"未找到帧号 {frame_num} 的对应行")  # 如果未找到匹配帧号，记录信息
            except Exception as e:
                logger.error(f"读取 TXT 文件失败: {e}")  # 记录读取文件失败的错误

        else:
            logger.warning(f"TXT 文件不存在: {txt_path}")  # 如果TXT文件不存在，记录警告

        # 将选择的行写入对应的choose.txt文件
        if choose_lines:
            choose_txt_name = txt_filename
            choose_txt_path = os.path.join(choose_output_dir, choose_txt_name)
            choose_txt_path = normalize_path(choose_txt_path)  # 确保路径分隔符与操作系统兼容

            # 以追加模式写入选择的行到对应的TXT文件
            with open(choose_txt_path, "a", encoding="utf-8") as ctf:
                for cline in choose_lines:
                    ctf.write(cline + "\n")

            # 将选择的行和对应的文件名写入汇总文件
            with open(choose_summary_path, "a", encoding="utf-8") as summary_f:
                for cline in choose_lines:
                    summary_f.write(f"{choose_txt_name}\t{cline}\n")

        # 将提取的目标对象转换为JSON字符串
        result_txt_data = json.dumps(extracted_objs, ensure_ascii=False)

        # 读取原始JSON文件内容
        try:
            with open(json_path, "r", encoding="utf-8") as jf:
                json_data = json.load(jf)
                json_str = json.dumps(json_data, ensure_ascii=False)  # 将JSON数据转换为字符串
        except Exception as e:
            json_str = f"读取JSON失败: {str(e)}"  # 如果读取失败，记录错误信息
            logger.error(f"读取 JSON 内容失败: {json_path}，错误：{e}")

        # 准备输出结果行
        out_line = f"{name_no_ext}\t{result_txt_data}\t{json_str}\n"

        # 获取父目录名称作为分组依据
        parent_dir = os.path.basename(os.path.dirname(json_path))
        processed_results[parent_dir].append(out_line)
        logger.debug(f"将结果添加到目录组: {parent_dir}")

# 按目录分组输出结果
empty_dirs = 0
valid_dirs = 0
for dir_name, result_lines in processed_results.items():
    # 检查该目录是否有有效的处理结果
    valid_results = [line for line in result_lines if '"cls":' in line and '"bbox":' in line]

    # 如果没有有效结果，跳过创建文件
    if not valid_results:
        logger.info(f"目录 {dir_name} 没有有效的匹配结果，跳过生成文件")
        empty_dirs += 1
        continue

    # 创建按目录命名的合并文件
    merge_file_name = f"{dir_name}_merge.txt"
    merge_file_path = os.path.join(merge_output_dir, merge_file_name)
    merge_file_path = normalize_path(merge_file_path)

    # 将该目录下所有处理结果写入到一个文件中
    with open(merge_file_path, "w", encoding="utf-8") as merge_f:
        for line in result_lines:
            merge_f.write(line)

    valid_dirs += 1
    logger.info(f"写入目录 {dir_name} 的合并文件完成：{merge_file_path}，共 {len(result_lines)} 条结果")

logger.info(f"处理完成，共处理 {file_count} 个文件，有效输出 {valid_dirs} 个合并文件，跳过 {empty_dirs} 个空目录")
logger.info(f"合并结果保存在: {merge_output_dir}")
logger.info(f"选择结果保存在: {choose_output_dir}")
